import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np
import sklearn
import sys
import os
import pandas as pd
import tensorflow as tf

from tensorflow.python import keras

print(tf.__version__)

root_directory = "../dataset/data4/"
train_directory = "train"
val_directory = "val"
height = 256
width = 256
channels = 3
batch_size = 64
num_classes = 2


# 数据增强
class LRN(keras.layers.Layer):
    def __init__(self):
        super(LRN, self).__init__()
        self.depth_radius = 2
        self.bias = 1
        self.alpha = 1e-4
        self.beta = 0.75

    def call(self, x):
        return tf.nn.lrn(x, depth_radius=self.depth_radius,
                         bias=self.bias, alpha=self.alpha,
                         beta=self.beta)


train_datagen = keras.preprocessing.image.ImageDataGenerator(
    rescale=1. / 255,
    rotation_range=15,  # 旋转40度
    width_shift_range=0.08,  # 位移 0.2 个比例 20% 平移
    height_shift_range=0.08,  #
    shear_range=0.1,  # 增强强度
    zoom_range=0.1,
    horizontal_flip=True,
    fill_mode='nearest',  # 填充
)
train_generator = train_datagen.flow_from_directory(root_directory+train_directory,
                                                    target_size=(height, width),
                                                    batch_size=batch_size,
                                                    seed=7,
                                                    shuffle=True,
                                                    class_mode="sparse")
validation_datagen = keras.preprocessing.image.ImageDataGenerator(
    rescale=1. / 255,
)
validation_generator = validation_datagen.flow_from_directory(root_directory+val_directory,
                                                              target_size=(height, width),
                                                              batch_size=batch_size,
                                                              seed=7,
                                                              shuffle=False,
                                                              class_mode="sparse")
train_num = train_generator.samples
valid_num = validation_generator.samples
print(train_num, valid_num)
# 建立模型
model = keras.models.Sequential()
model.add(keras.layers.Conv2D(filters=96,
                              kernel_size=(11, 11),
                              strides=4,
                              activation='relu',
                              padding='same',
                              input_shape=[width, height, channels]))
model.add(keras.layers.MaxPool2D(pool_size=(3, 3), strides=2))
model.add(LRN())
model.add(keras.layers.Conv2D(filters=256,kernel_size=(5, 5),strides=1,activation='relu',padding='same'))
model.add(keras.layers.MaxPool2D(pool_size=(3, 3), strides=2))
model.add(LRN())
model.add(keras.layers.Conv2D(filters=384, kernel_size=(3, 3),strides=1,activation='relu',padding='same'))

model.add(keras.layers.Conv2D(filters=384,kernel_size=(3, 3),strides=1,activation='relu',padding='same'))
model.add(keras.layers.Conv2D(filters=256,kernel_size=(3, 3),strides=1,activation='relu',padding='same'))
model.add(keras.layers.MaxPool2D(pool_size=(3, 3), strides=2))

model.add(keras.layers.Flatten())
model.add(keras.layers.Dense(4096, activation='relu'))
model.add(keras.layers.Dropout(0.5))
model.add(keras.layers.Dense(4096, activation='relu'))
model.add(keras.layers.Dropout(0.5))
model.add(keras.layers.Dense(num_classes, activation="softmax"))
print(model.summary())

model.compile(loss="categorical_crossentropy",
              optimizer="adam",
              metrics=["accuracy"])
epochs = 300

history = model.fit_generator(train_generator,
                              steps_per_epoch=train_num // batch_size,
                              epochs=epochs,
                              validation_data=validation_generator,
                              validation_steps=valid_num // batch_size)
model.save_weights("./model/alexNet_300",save_format="tf")

def plot_learing_curves(history, label, epochs, min_value, max_value):
    data = {}
    data[label] = history.history[label]
    data['val' + label] = history.history['val_' + label]
    pd.DataFrame(data).plot(figsize=(8, 5))
    plt.grid(True)
    plt.axis([0, epochs, min_value, max_value])
    plt.show()


plot_learing_curves(history, 'accuracy', epochs, 0, 1)
plot_learing_curves(history, 'loss', epochs, 0, 1)
